#include <TrackingThread.hxx>
#include "FramePool.hxx"
#include "Sequence.hxx"
#include "VideoThread.hxx"
#include "TrackedObject.hxx"



TrackingThread::TrackingThread()
{
	_frameBuffer = 0;
	_stopTrackFlag = false;
	_trackingFlag = false;
	_pauseFlag = false;
	_trackerInitedFlag = false;
	_slowFlag = false;
	_nextFlag = false;
	_stopFlag = false;

	strcpy(_trackerType, "CSRT");

	_status = TR_STATUS_FREE;
}

TrackingThread::~TrackingThread()
{
	if (_frameBuffer)
	{
		delete[] _frameBuffer;
		_frameBuffer = 0;
	}
}

void TrackingThread::initialize()
{
	int pageSize = _inputPool->getPageSize();
	_outputPool->setPageSize(pageSize);
	_outputPool->setPageCount(10);
	_outputPool->initialize();

	_frameBuffer = new char[pageSize];
}

void TrackingThread::run()
{
	while (1)
	{
		if (_stopFlag) break;

		if (_pauseFlag)
		{
			msleep(30);
			continue;
		}

		if (_slowFlag)
			msleep(100);

		if (_inputPool->popFrame(_frameBuffer))
		{
			int frameIndex = *(int*)_frameBuffer;

			static int lastFrameIndex = frameIndex;

			if ((frameIndex < lastFrameIndex) || (frameIndex - lastFrameIndex != 1))
			{
				printf("Error!!!\n");
			}
			lastFrameIndex = frameIndex;

			if (_trackingFlag)
			{
				cv::Mat img(_model->getVideoHeight(), _model->getVideoWidth(), CV_8UC3, _frameBuffer + sizeof(VideoFrameInfo));

				static bool trackerInited = false;
				cv::Rect2d bbox;
				std::vector<TrackedObject*>* objects = _model->getTrackingObjects();
				if (!_trackerInitedFlag)
				{					
					_tracker = cv::MultiTracker::create();
					for (int i = 0; i < objects->size(); i++)
					{
						TrackedObject* tobj = objects->at(i);
						
						// add template frame to tracked object
						tobj->setStartFrameIndex(frameIndex);
						tobj->setEndFrameIndex(frameIndex);
						tobj->pushRect(frameIndex, tobj->getTemplateBox());
						
						if( strcmp(_trackerType, "CSRT") == 0 )
							_tracker->add(cv::TrackerCSRT::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "MOSSE") == 0)
							_tracker->add(cv::TrackerMOSSE::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "GOTURN") == 0)
							_tracker->add(cv::TrackerGOTURN::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "KCF") == 0)
							_tracker->add(cv::TrackerKCF::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "TLD") == 0)
							_tracker->add(cv::TrackerTLD::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "MedianFlow") == 0)
							_tracker->add(cv::TrackerMedianFlow::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "Boosting") == 0)
							_tracker->add(cv::TrackerBoosting::create(), img, tobj->getTemplateBox());
						else if (strcmp(_trackerType, "MIL") == 0)
							_tracker->add(cv::TrackerMIL::create(), img, tobj->getTemplateBox());
					}
					_trackerInitedFlag = true;

					_status = TR_STATUS_RUNNING;
				}
				else
				{
					_tracker->update(img);

					if (_stopTrackFlag)
					{
						_status = TR_STATUS_STOPPED;
						_trackingFlag = false;
						_stopTrackFlag = false;
					}
				}

				for (unsigned i = 0; i<_tracker->getObjects().size(); i++)
				{
					
					TrackedObject* tobj = objects->at(i);
					cv::Rect2d rect = _tracker->getObjects()[i];
					tobj->setEndFrameIndex(frameIndex);
					tobj->pushRect(frameIndex, rect);
					//cv::rectangle(img, , cv::Scalar(255, 255, 0), 3);
				}
			}

			while (_outputPool->isFull())
				msleep(30);

			if (!_outputPool->pushFrame(_frameBuffer))
			{
				msleep(20);
			}
		}
		else
		{
			msleep(20);
		}


		if (_nextFlag)
		{
			_pauseFlag = true;
			_nextFlag = false;
		}
	}
}

void TrackingThread::startTrack()
{
	_trackingFlag = true;
}

void TrackingThread::stopTrack()
{
	_stopTrackFlag = true;
	
	while( this->isRunning() && ( _status == TR_STATUS_RUNNING ))
		msleep(500);

	_tracker.release();
	
	_trackerInitedFlag = false;
}

void TrackingThread::pause()
{
	_pauseFlag = true;
}

void TrackingThread::resume()
{
	_pauseFlag = false;
}

void TrackingThread::slow()
{
	_slowFlag = true;
}

void TrackingThread::fast()
{
	_slowFlag = false;
}

void TrackingThread::stop()
{
	_stopFlag = true;
}

void TrackingThread::next()
{
	_nextFlag = true;
	_pauseFlag = false;
}
